import asyncio
import threading
import traceback
import weakref
from concurrent.futures import ProcessPoolExecutor
from concurrent.futures.process import BrokenProcessPool
from functools import partial
from typing import Callable, List, Optional

from areal.utils import logging

logger = logging.getLogger("Reward API")


def reward_fn(
    prompt: str,
    completions: str,
    prompt_ids: List[int],
    completion_ids: List[int],
    **kwargs,
):
    """This function is a placeholder for the reward function that will be used in the RLVR pipeline.

    In general, there's no restriction on the signature and implementation of this function in customized rollout workflows.
    It would be convinent to follow this signature and directly use it in our predefined rollout workflows.

    :param prompt: The string representing the task to be completed.
    :param completions: The string representing the trajectory generated by the model.
    :param prompt_ids: The token IDs of the prompt.
    :param completion_ids: The token IDs of the trajectory generated by the model.
    :param kwargs: Other attributes of the data in the dataset, such as solutions, input_outputs, etc.
        Any other attributes in the dataset will be passed as keyword arguments to this function.
    :rtype: float
    """


class AsyncRewardWrapper:
    """
    Wraps a synchronous reward function to make it async with timeout handling.
    Automatically manages ProcessPoolExecutor lifecycle based on instance count.
    Includes automatic recovery from broken process pools.
    """

    _executors = {}
    _instance_counts = {}
    _lock = threading.Lock()

    def __init__(
        self,
        reward_fn: Callable,
        timeout_seconds: float = 15,
        max_workers: Optional[int] = None,
        max_retries: int = 3,
    ):
        self.reward_fn = reward_fn
        self.timeout_seconds = timeout_seconds
        self.max_workers = max_workers
        self.max_retries = max_retries
        self._executor_key = max_workers

        with self._lock:
            if self._executor_key not in self._executors:
                self._executors[self._executor_key] = ProcessPoolExecutor(
                    max_workers=max_workers
                )
                self._instance_counts[self._executor_key] = 0
            self._instance_counts[self._executor_key] += 1

        weakref.finalize(self, AsyncRewardWrapper._cleanup_executor, max_workers)

    @classmethod
    def _cleanup_executor(cls, executor_key):
        """Called when an AsyncRewardWrapper instance is garbage collected"""
        with cls._lock:
            if executor_key in cls._instance_counts:
                cls._instance_counts[executor_key] -= 1
                if cls._instance_counts[executor_key] <= 0:
                    if executor_key in cls._executors:
                        executor = cls._executors.pop(executor_key)
                        executor.shutdown(wait=True)
                        logger.debug(
                            f"ProcessPoolExecutor with {executor_key} workers shut down"
                        )
                    cls._instance_counts.pop(executor_key, None)

    @classmethod
    def _recreate_executor(cls, executor_key, max_workers):
        """Recreate a broken ProcessPoolExecutor"""
        with cls._lock:
            if executor_key in cls._executors:
                # Clean up the broken executor
                old_executor = cls._executors[executor_key]
                try:
                    old_executor.shutdown(wait=False)
                except Exception as e:
                    logger.warning(f"Error shutting down broken executor: {e}")

                # Create a new executor
                cls._executors[executor_key] = ProcessPoolExecutor(
                    max_workers=max_workers
                )
                logger.info(f"Recreated ProcessPoolExecutor with {max_workers} workers")
                return cls._executors[executor_key]
        return None

    async def __call__(self, *args, **kwargs) -> float:
        last_exception = None

        for attempt in range(self.max_retries + 1):
            with self._lock:
                executor = self._executors.get(self._executor_key)

            if executor is None:
                raise RuntimeError("ProcessPoolExecutor has been shut down")

            loop = asyncio.get_event_loop()
            try:
                return await asyncio.wait_for(
                    loop.run_in_executor(
                        executor,
                        partial(self.reward_fn, *args, **kwargs),
                    ),
                    timeout=self.timeout_seconds,
                )
            except asyncio.TimeoutError:
                logger.warning(
                    f"Computing reward timeout after {self.timeout_seconds}s. Set reward to 0."
                )
                return 0
            except BrokenProcessPool as e:
                last_exception = e
                logger.warning(
                    f"ProcessPoolExecutor broken (attempt {attempt + 1}/{self.max_retries + 1}). "
                    "Attempting to recreate..."
                )
                if attempt < self.max_retries:
                    # Try to recreate the executor
                    new_executor = self._recreate_executor(
                        self._executor_key, self.max_workers
                    )
                    if new_executor is None:
                        logger.error("Failed to recreate ProcessPoolExecutor")
                        break
                    # Continue to next attempt
                    continue
                else:
                    logger.error("Max retries exceeded for BrokenProcessPool.")
                    traceback.print_exc()
                    raise e
            except Exception as e:
                last_exception = e
                logger.error(f"Unexpected error in reward computation: {e}")
                if attempt < self.max_retries:
                    logger.info(
                        f"Retrying... (attempt {attempt + 1}/{self.max_retries + 1})"
                    )
                    continue
                else:
                    logger.error("Max retries exceeded for unexpected error.")
                    traceback.print_exc()
                    raise e

        # If we get here, all retries failed
        if last_exception:
            traceback.print_exc()
            raise last_exception
        else:
            raise RuntimeError("Reward computation failed after all retries.")
